'''
后面几个提取器接收的都是packet列表
直接减少多次读取同一文件的开销

本文件的功能是提取流中的包序列特征
序列方向定义：从主机向外发送的包序列为正向+1，从外向主机发送的包序列为负向-1

提取信息包含：
包序列
正序列长度的 最大值 最小值 平均值 标准差
负序列长度的 最大值 最小值 平均值 标准差
'''
'''
外部调用接口
生成实例对象，传入packet列表，调用Extract_Features方法，即可得到流的特征
值可以通过print打印出来，也可以直接获取写入
'''

from scapy.all import *
import numpy as np

class PacketFeatureExtractor:
    def __init__(self, packet_list, dfmodel, hostIP = "127.0.0.1"):
        self.dfmodel = dfmodel
        self.hostIP = hostIP
        self.packet_list = packet_list 
        self.packet_sequence = []
        
        self.forward_max = 0
        self.forward_min = 0
        self.forward_mean = 0
        self.forward_std = 0
        
        self.backward_max = 0    
        self.backward_min = 0
        self.backward_mean = 0
        self.backward_std = 0
        
        self.Extract_Features()

    def Extract_Features(self):
        for packet in self.packet_list:
            if packet["IP"].src == self.hostIP:
                self.packet_sequence.append(1)
            else:
                self.packet_sequence.append(-1)


        forward_sequence = []
        backward_sequence = []
        
        # a作为比较位， b作为计数器
        a = self.packet_sequence[0]
        b = 0
        for x in self.packet_sequence:
            if x == a:
                b += 1
            else:
                if a > 0:
                    forward_sequence.append(b)
                else:
                    backward_sequence.append(b)
                a = x
                b = 1
                
            if a > 0:
                forward_sequence.append(b)
            else:    
                backward_sequence.append(b)
        
        # 防止出现空列表导致运行暂停
        try:
            self.forward_max = max(forward_sequence)
            self.forward_min = min(forward_sequence)
            self.forward_mean = np.mean(forward_sequence)
            self.forward_std = np.std(forward_sequence)
        except:
            pass
    
        try:
            self.backward_max = max(backward_sequence)
            self.backward_min = min(backward_sequence)
            self.backward_mean = np.mean(backward_sequence)
            self.backward_std = np.std(backward_sequence)
        except:
            pass
        
        # 使用模型处理包特征，需要在最后进行
        # 输入必须是一个二维列表
        self.packet_sequence = self.UtilizeDFmodel([self.packet_sequence])
    
    def UtilizeDFmodel(self,packet_sequence):
    # 使用训练好的DF模型进行预测
    
        # 将列表裁剪为5000长度，转换为np.array
        for item in packet_sequence:
            if len(item) > 5000:
                item = item[:5000]
            elif len(item) < 5000:
                item.extend([0] * (5000 - len(item)))

        v = np.array(packet_sequence)

        # 输入的应该为三维数组
        v = v.astype('float32')
        v = v[:, :,np.newaxis]

        layer_out = self.dfmodel.predict(v)
        
        # print("The length of layer_out is: ", len(layer_out[0]))
        return layer_out[0]
    
    
    def __str__(self):
        return "Packet Sequence: " + str(self.packet_sequence) + "\n" + \
                "Forward Max: " + str(self.forward_max) + "\n" + \
                "Forward Min: " + str(self.forward_min) + "\n" + \
                "Forward Mean: " + str(self.forward_mean) + "\n" + \
                "Forward Std: " + str(self.forward_std) + "\n" + \
                "Backward Max: " + str(self.backward_max) + "\n" + \
                "Backward Min: " + str(self.backward_min) + "\n" + \
                "Backward Mean: " + str(self.backward_mean) + "\n" + \
                "Backward Std: " + str(self.backward_std) + "\n"
                
    def Return_Features(self):
        # 返回特征值列表
        return list(self.packet_sequence) + [self.forward_max, self.forward_min, self.forward_mean, self.forward_std, self.backward_max, self.backward_min, self.backward_mean, self.backward_std]
    
if __name__ == "__main__":
    packet_list = rdpcap("./processed_dataset/browsing_flow_4.pcap")
    extractor = PacketFeatureExtractor(packet_list, hostIP="10.152.152.11")
    print(extractor)